from datetime import datetime
from typing import List

from elasticsearch import Elasticsearch
from langchain_community.document_compressors import JinaRerank
from langchain_community.embeddings import JinaEmbeddings
from langchain_core.documents import Document
from openai import OpenAI

from es_retriever import get_relevant_documents

# 存储到Elasticsearch
embeddings = JinaEmbeddings(
    jina_api_key="jina_21eda124ed7e4ee895f7bb7197f3cc5ebT45ebFpQaWsIeieL6QYTY9Sojgu",
    model_name="jina-embeddings-v2-base-zh"
)
rag_index = "rag_documents"
file_index = "rag_file_info"
es_url = "http://10.252.202.55:9200"  # Elasticsearch地址

es_client = Elasticsearch(hosts=[es_url])
print(es_client.info())
user_id = 10086
cur_query = '项目'
print("当前的查询:" + cur_query)

index_name = rag_index
text_field = "text"
dense_vector_field = "vector"


def hybrid_query(search_query: str) -> List[Document]:
    # 在查询时动态添加过滤条件
    results = get_relevant_documents(search_query, es_client=es_client,
                                     index_name=rag_index,
                                     embedding_model=embeddings,
                                     keyword_boost=0.5,
                                     vector_boost=0.5,
                                     default_filters=[
                                         {"term": {"metadata.user_id": user_id}}  # 默认过滤条件
                                     ])
    print(results)
    return results


docs = hybrid_query(cur_query)

compressor = JinaRerank(jina_api_key="jina_21eda124ed7e4ee895f7bb7197f3cc5ebT45ebFpQaWsIeieL6QYTY9Sojgu",
                        model="jina-reranker-v2-base-multilingual")

re_docs = compressor.rerank(query=cur_query, documents=docs, top_n=5)

print(re_docs)

final_docs = []
for doc in re_docs:
    index = doc['index']
    final_docs.append(docs[index])

print(final_docs)

# 配置DeepSeek API
client = OpenAI(
    base_url='https://api-inference.modelscope.cn/v1/',
    api_key='b011f8fd-99a5-4633-8bf5-fef46cf7bd90',  # ModelScope Token
)


def build_system_prompt(search_results: str, query: str) -> str:
    """构建系统提示模板"""
    current_time = datetime.now().strftime('%Y-%m-%d %H:%M:%S')

    return f"""
# 以下内容是基于用户发送的消息的搜索结果:{search_results}

在我给你的搜索结果中，每个结果都是[chunk X begin]...[chunk X end]格式的，X代表每篇文章某个切片的数字索引，一个文章可能包含多个切片，顺序是乱序。
请在适当的情况下在句子末尾引用上下文。请按照引用编号[citation:X]的格式在答案中对应部分引用上下文。如果一句话源自多个上下文，请列出所有相关的引用编号，例如[citation:3][citation:5]，切记不要将引用集中在最后返回引用编号，而是在答案对应部分列出。
在回答时，请注意以下几点：
- 今天是{current_time}。
- 并非搜索结果的所有内容都与用户的问题密切相关，你需要结合问题，对搜索结果进行甄别、筛选。
- 对于列举类的问题（如列举所有航班信息），尽量将答案控制在10个要点以内，并告诉用户可以查看搜索来源、获得完整信息。优先提供信息完整、最相关的列举项；如非必要，不要主动告诉用户搜索结果未提供的内容。
- 对于创作类的问题（如写论文），请务必在正文的段落中引用对应的参考编号，例如[citation:3][citation:5]，不能只在文章末尾引用。你需要解读并概括用户的题目要求，选择合适的格式，充分利用搜索结果并抽取重要信息，生成符合用户要求、极具思想深度、富有创造力与专业性的答案。你的创作篇幅需要尽可能延长，对于每一个要点的论述要推测用户的意图，给出尽可能多角度的回答要点，且务必信息量大、论述详尽。
- 如果回答很长，请尽量结构化、分段落总结。如果需要分点作答，尽量控制在5个点以内，并合并相关的内容。
- 对于客观类的问答，如果问题的答案非常简短，可以适当补充一到两句相关信息，以丰富内容。
- 你需要根据用户要求和回答内容选择合适、美观的回答格式，确保可读性强。
- 你的回答应该综合多个文件来回答，不能重复引用一个文件。
- 除非用户要求，否则你回答的语言需要和用户提问的语言保持一致。
# 用户消息为：{query}"""


# 按doc_id分组文档





# merged_docs现在包含合并后的文档列表

# TODO 需要将同一个文章的内容合并到一起，按照chun_id
formatted_result = ""
for i, item in enumerate(final_docs, start=1):
    chunk_id = item.metadata['chunk_id']
    filename = item.metadata['filename']
    page= item.page_content[:2000]
    formatted_result += (
        f"[chunk {chunk_id} begin]\n"
        f"filename: {filename}\n"
        f"content: {page}...\n"  # 限制内容长度
        f"[chunk {chunk_id} end]\n"
    )

user_prompt = build_system_prompt(formatted_result, cur_query)

print(user_prompt)

response = client.chat.completions.create(
    model='deepseek-ai/DeepSeek-R1-Distill-Qwen-7B',
    messages=[
        {"role": "user", "content": user_prompt}
    ],
    stream=True,
    temperature=0.7,  # 降低随机性
    max_tokens=8000
)

print("\n\n=== 生成回答 ===")
done_reasoning = False
for chunk in response:
    reasoning_chunk = chunk.choices[0].delta.reasoning_content
    answer_chunk = chunk.choices[0].delta.content
    if reasoning_chunk != '':
        print(reasoning_chunk, end='', flush=True)
    elif answer_chunk != '':
        if not done_reasoning:
            print('\n\n === Final Answer ===\n')
            done_reasoning = True
        print(answer_chunk, end='', flush=True)

print(222)



def merge_doc():
    doc_groups = {}
    for item in final_docs:
        doc_id = item.metadata['doc_id']
        if doc_id not in doc_groups:
            doc_groups[doc_id] = []
        doc_groups[doc_id].append(item)
    # 对每个分组排序并合并内容
    merged_docs = []
    for doc_id, items in doc_groups.items():
        # 按chunk_id排序（转换为整数确保正确排序）
        sorted_items = sorted(items, key=lambda x: int(x.metadata['chunk_id']))

        # 合并内容（每个chunk取前2000字符）
        merged_content = ''.join([item.page_content for item in sorted_items])

        # 获取公共元数据（假设同doc_id的文件名一致）
        metadata = {
            'doc_id': doc_id,
            'filename': sorted_items[0].metadata['filename'],
            'total_chunks': len(sorted_items)
        }

        # 构造合并后的文档对象（可根据需要调整数据结构）
        merged_docs.append({
            **metadata,
            'content': merged_content
        })
